import os
import pickle

import numpy as np
import tensorflow.compat.v1 as tf

from utils.tf_record import build_int64_feature, build_bytes_feature, build_example

classification = ['airplane',
                  'automobile',
                  'bird',
                  'cat',
                  'deer',
                  'dog',
                  'frog',
                  'horse',
                  'ship',
                  'truck']


def unpickle(filename):
    """Decode the dataset files."""
    with open(filename, 'rb') as f:
        d = pickle.load(f, encoding='latin1')
        return d


def load_data(dataset_dir, train_filenames, test_filenames):
    train_images = unpickle(os.path.join(dataset_dir, train_filenames[0]))['data']
    train_labels = unpickle(os.path.join(dataset_dir, train_filenames[0]))['labels']
    test_images = unpickle(os.path.join(dataset_dir, test_filenames[0]))['data']
    test_labels = unpickle(os.path.join(dataset_dir, test_filenames[0]))['labels']
    for i in range(1, len(train_filenames)):
        batch = unpickle(os.path.join(dataset_dir, train_filenames[i]))
        train_images = np.concatenate((train_images, batch['data']), axis=0)
        train_labels = np.concatenate((train_labels, batch['labels']), axis=0)
    return train_images, train_labels, test_images, test_labels


def save_to_record(images, labels, data_type='train'):
    assert len(images) == len(labels)
    with tf.io.TFRecordWriter('../data/' + data_type + '.tfrecord') as w:
        for i in range(len(images)):
            feature = {
                "image": build_bytes_feature(images[i].tostring()),
                "label": build_int64_feature(labels[i]),
            }
            example = build_example(feature)
            w.write(example.SerializeToString())


def main():
    dataset_dir = '../data/cifar-10-batches-py'
    train_filenames = [
        'data_batch_1',
        'data_batch_2',
        'data_batch_3',
        'data_batch_4',
        'data_batch_5'
    ]
    test_filenames = [
        'test_batch'
    ]
    train_images, train_labels, test_images, test_labels = load_data(dataset_dir, train_filenames, test_filenames)
    save_to_record(train_images, train_labels, data_type='train')
    save_to_record(test_images, test_labels, data_type='test')


if __name__ == '__main__':
    main()
